package tools

import (
	"context"
	"errors"
	"fmt"
	"sync"
)

type ProduceConsumerQueue[T any] struct {
	dataChannel chan T
	consumers   []*consumer[T]
	ctx         context.Context
	lock        sync.RWMutex
	stop        bool
}

func NewProduceConsumerQueue[T any](ctx context.Context, channelLength int) *ProduceConsumerQueue[T] {
	if ctx == nil {
		ctx = context.Background()
	}
	if channelLength < 0 {
		channelLength = 0
	}
	queue := &ProduceConsumerQueue[T]{
		dataChannel: make(chan T, channelLength),
		consumers:   make([]*consumer[T], 0),
		ctx:         ctx,
	}
	go queue.startConsume()
	return queue
}

func (q *ProduceConsumerQueue[T]) Produce(data T) error {
	if q.stop {
		return errors.New("queue is already stopped")
	}
	go func(sendData T) {
		q.dataChannel <- sendData
	}(data)
	return nil
}

// RegisterAsyncConsumer 注册异步消费函数
func (q *ProduceConsumerQueue[T]) RegisterAsyncConsumer(consumerFun func(T)) error {
	if consumerFun == nil {
		return errors.New("consumerFun is nil")
	}
	q.lock.Lock()
	q.consumers = append(q.consumers, &consumer[T]{
		async:      true,
		consumeFun: consumerFun,
	})
	q.lock.Unlock()
	return nil
}

// RegisterSyncConsumer 注册同步消费函数
func (q *ProduceConsumerQueue[T]) RegisterSyncConsumer(consumerFun func(T)) error {
	if consumerFun == nil {
		return errors.New("consumerFun is nil")
	}
	q.lock.Lock()
	q.consumers = append(q.consumers, &consumer[T]{
		async:      false,
		consumeFun: consumerFun,
	})
	q.lock.Unlock()
	return nil
}

func (q *ProduceConsumerQueue[T]) startConsume() {
	for {
		select {
		case <-q.ctx.Done():
			fmt.Println("queue stop")
			q.stop = true
			return
		case data := <-q.dataChannel:
			q.lock.RLock()
			for _, tempConsumer := range q.consumers {
				consumerStruct := tempConsumer
				if consumerStruct.async {
					go invokeConsumeFunc(consumerStruct.consumeFun, data)
				} else {
					invokeConsumeFunc(consumerStruct.consumeFun, data)
				}
			}
			q.lock.RUnlock()
		}
	}
}

func invokeConsumeFunc[T any](consumeFun func(T), data T) {
	defer func() {
		// 忽略调用错误
		_ = recover()
	}()
	if consumeFun == nil {
		return
	}
	consumeFun(data)
}

type consumer[T any] struct {
	async      bool
	consumeFun func(T)
}
